import torch
import torch.nn as nn
import os
import torch.nn.functional as F
import gradio as gr
from PIL import Image
from torchvision import datasets, transforms
model_path = "best_mlp.pth"

class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)  # Flatten the input
        x = F.relu(self.fc1(x))
        x = F.dropout(x, p=0.2, training=self.training)
        x = F.relu(self.fc2(x))
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.fc3(x)
        return F.log_softmax(x, dim=1)

model_size = os.path.getsize(model_path)  # 获取完整模型文件的大小，用于打印文件信息
print(f"模型参数文件路径: {model_path}")  # 打印模型参数文件路径
print(f"模型参数文件大小: {model_size} bytes")  # 打印模型参数文件大小

model_params = torch.load(model_path)
model = MLP()
model.load_state_dict(model_params)
model.eval()

print("载入的模型参数:")  # 打印提示信息，表示开始打印载入的模型参数
for name, param in model.named_parameters():  # 遍历模型的所有参数
    print(f"{name}: {param.size()}")  # 打印参数名称和形状

print("使用模型参数载入的模型结构信息:")  # 打印提示信息，表示开始打印载入的模型结构信息
print(model)


# 定义预处理函数
def preprocess(image):
    image = Image.fromarray(image).convert('L')  # 转换为灰度图像
    image_tensor = transforms.ToTensor()(image)  # 转换为张量
    image_tensor = image_tensor.view(1, -1)  # 添加批次维度
    return image_tensor

# 定义预测函数
def predict(image):
    preprocessed_image = preprocess(image)
    model.eval()  # Set the model to evaluation mode
    with torch.no_grad():
        output = model(preprocessed_image)
        _, predicted_digit = torch.max(output, 1)
    print(predicted_digit.item())
    return predicted_digit.item()

# 创建Gradio接口
iface = gr.Interface(fn=predict, inputs='sketchpad', outputs='text')

# 启动Gradio接口
iface.launch()







